import numpy as np
import matplotlib.pyplot as plt

## GWO优化算法
def gwo(agents_num, max_iter, dim, lb, ub):
    '''
    agents_num: 狼群数量
    max_iter: 最大迭代次数
    dim: 需要优化参数个数
    lb： 参数取值下界
    ub： 参数取值上界
    '''
    # 初始化Alpha狼的位置
    Alpha_pos=[0,0]
    Beta_pos=[0,0]
    Delta_pos=[0,0]
    # 初始化Alpha狼的目标函数值
    Alpha_score = float("inf")
    Beta_score = float("inf")
    Delta_score = float("inf")
    # 初始化首次搜索位置
    pos = np.dot(np.random.rand(agents_num, dim), (ub - lb)) + lb
    iterations = []
    accuracy = []
    # 循环次数
    iter = 0
    while iter < max_iter:
        # t0 = time()
        # 遍历每个狼
        for i in range(0, (pos.shape[0])):
            # 若搜索位置超过了搜索空间，需要重新回到搜索空间 
            for j in range(0, (pos.shape[1])):
                Flag4ub = pos[i, j] > ub
                Flag4lb = pos[i, j] < lb
                # 若狼的位置在最大值和最小值之间，则位置不需要调整，若超出最大值，最回到最大值边界
                if Flag4ub:
                    pos[i, j] = ub
                if Flag4lb:
                    pos[i, j] = lb
            fitness = 0
            # LSSVM
            # sigma, alpha, b = lssvm_fit(x_train, y_train, C=pos[i][0], sigma=pos[i][1])
            # y_pre = lssvm_predict(sigma, alpha, b, x_train, x_train)
            # fitness = r2_score(y_train, y_pre) * 100
            #如果目标函数值小于Alpha狼的目标函数值
            if fitness < Alpha_score:
                Alpha_score = fitness # 则将Alpha狼的目标函数值更新为最优目标函数值
                Alpha_pos = pos[i] #同时将Alpha狼的位置更新为最优位置
            #如果目标函数值介于于Alpha狼和Beta狼的目标函数值之间
            if fitness > Alpha_score and fitness < Beta_score:
                Beta_score = fitness # 则将Beta狼的目标函数值更新为最优目标函数值
                Beta_pos = pos[i]
            #如果目标函数值介于于Beta狼和Delta狼的目标函数值之间
            if fitness > Alpha_score and fitness > Beta_score and fitness < Delta_score:
                Delta_score = fitness # 则将Delta狼的目标函数值更新为最优目标函数值
                Delta_pos = pos[i]
        a = 2 - iter * (2 / max_iter)
        
        # 遍历每个狼
        for i in range(0, (pos.shape[0])):
            #遍历每个维度
            for j in range(0, (pos.shape[1])):
                #包围猎物，位置更新                
                r1 = np.random.random(1)#生成0~1之间的随机数
                r2 = np.random.random(1)
                A1 = 2 * a * r1 - a # 计算系数A
                C1 = 2 * r2 # 计算系数C
                #Alpha狼位置更新
                D_alpha = abs(C1 * Alpha_pos[j] - pos[i,j])
                X1 = Alpha_pos[j] - A1 * D_alpha
                r1 = np.random.random(1)
                r2 = np.random.random(1)
                A2 = 2 * a * r1 - a
                C2 = 2 * r2
                # Beta狼位置更新
                D_beta = abs(C2 * Beta_pos[j] - pos[i,j])
                X2 = Beta_pos[j] - A2 * D_beta
                r1 = np.random.random(1)
                r2 = np.random.random(1)
                A3 = 2 * a * r1 - a
                C3 = 2 * r2
                # Delta狼位置更新
                D_delta = abs(C3 * Delta_pos[j] - pos[i,j])
                X3 = Delta_pos[j] - A3 * D_delta
                # 位置更新
                pos[i,j] = (X1 + X2 + X3) / 3
        iter += 1
        iterations.append(iter)
        accuracy.append((100 - Alpha_score) / 100)
        # print(f'迭代次数: {iter}, {(time() - t0):.0f}s')
    # 绘制迭代曲线
    _ = plt.plot(iterations, accuracy)
    _ = plt.xlabel('迭代次数', size = 20)
    _ = plt.ylabel('精度', size = 20)
    _ = plt.title('GWO-LSSVM参数优化')
    _ = plt.show()
    bestC = Alpha_pos[0]
    bestsigma = Alpha_pos[1]
    return bestC, bestsigma